#include "le_core.h"
#include "exr_decode_example_app.h"

#include "le_window.h"
#include "le_renderer.hpp"
#include "le_camera.h"
#include "le_pipeline_builder.h"

#include "le_mesh_generator.h"
#include "le_mesh.h"
#include "le_resource_manager.h"
#include "le_exr.h"
#include "le_ui_event.h"

#include "glm/glm.hpp"
#include "glm/gtc/matrix_transform.hpp"

#include <iostream>
#include <memory>
#include <sstream>
#include <vector>

#include <cassert>

struct GpuMeshData {
	le_buffer_resource_handle pos_handle;
	le_buffer_resource_handle uv_handle;
	le_buffer_resource_handle index_handle;
	size_t                 pos_count;
	size_t                 uv_count;
	size_t                 indices_count;
	size_t                 pos_num_bytes;
	size_t                 uv_num_bytes;
	size_t                 indices_num_bytes;
	le::IndexType          index_type;
};

struct exr_decode_example_app_o {
	le::Window   window;
	le::Renderer renderer;

	uint64_t frame_counter = 0;

	GpuMeshData* gpu_mesh = nullptr; // owning

	LeCamera                 camera;
	LeCameraController       cameraController;
	LeResourceManager        resource_manager{ renderer };
	le_image_resource_handle img_heightmap = renderer.createImageResourceHandle( "heightmap_image" );
	le_texture_handle        tex_unit_0    = renderer.produceTextureHandle( "tex_unit_0" );
	LeMesh                   mesh;
	bool                     was_mesh_uploaded = false;
};

static void reset_camera( exr_decode_example_app_o* self );                             // ffdecl.
static void exr_decode_example_app_process_ui_events( exr_decode_example_app_o* self ); // ffdecl.

// ----------------------------------------------------------------------

static void app_initialize() {
	le::Window::init();
};

// ----------------------------------------------------------------------

static void app_terminate() {
	le::Window::terminate();
};

// ----------------------------------------------------------------------

static exr_decode_example_app_o* exr_decode_example_app_create() {
	auto app = new ( exr_decode_example_app_o );

	// Register .exr image decoder so that resource manager knows what to do with exr files
	app->resource_manager.set_decoder_interface_for_filetype( "exr", le_exr_api_i->le_exr_image_decoder_i );

	le::Window::Settings settings;
	settings
	    .setWidth( 1024 )
	    .setHeight( 1024 )
	    .setTitle( "Island // ExrDecodeExampleApp" );

	// Create a new window
	app->window.setup( settings );

	// Set up the renderer
	app->renderer.setup( app->window );

	{
		LeMeshGenerator::generatePlane( app->mesh, 1024, 1024, 256, 256 );

		uint32_t index_num_bytes_per_index = 0;
		size_t   index_count               = 0;

		index_count = le_mesh::le_mesh_i.get_index_count( app->mesh, &index_num_bytes_per_index );

		size_t vertex_count = le_mesh::le_mesh_i.get_vertex_count( app->mesh );

		le_mesh_api::attribute_info_t attribute_info[ 4 ] = {};
		size_t                        num_attributes      = sizeof( attribute_info ) / sizeof( attribute_info[ 0 ] );

		le_mesh::le_mesh_i.read_attribute_infos_into( app->mesh, attribute_info, &num_attributes );

		auto get_num_bytes = [ & ]( le_mesh_api::attribute_name_t name ) -> size_t {
			for ( auto& e : attribute_info ) {
				if ( e.name == name ) {
					return e.bytes_per_vertex;
				}
			}
			return 0;
		};

		// Initialize handles, and size infos for gpu mesh data:
		app->gpu_mesh = new GpuMeshData{
		    app->renderer.createBufferResourceHandle( "vertex_buffer" ),
		    app->renderer.createBufferResourceHandle( "uv_buffer" ),
		    app->renderer.createBufferResourceHandle( "index_buffer" ),
		    vertex_count,
		    vertex_count,
		    index_count,
		    vertex_count * get_num_bytes( le_mesh_api::ePosition ),
		    vertex_count * get_num_bytes( le_mesh_api::eUv ),
		    index_count * index_num_bytes_per_index,
		    index_num_bytes_per_index == 2 ? le::IndexType::eUint16 : le::IndexType::eUint32,
		};
	}

	auto window_extents = app->renderer.getSwapchainExtent();
	app->cameraController.setControlRect( 0, 0, float( window_extents.width ), float( window_extents.height ) );
	// Set up the camera
	reset_camera( app );

	// We will load a grayscale '.exr' image with 32bit float pixels --
	// we must provide a GPU image that has matching specs:

	auto image_info =
	    le::ImageResourceInfoBuilder()
	        .setFormat( le::Format::eR32Sfloat )
	        .setUsageFlags( le::ImageUsageFlagBits::eSampled | le::ImageUsageFlagBits::eTransferDst )
	        .build();

	// Height map generated by hand
	// char const* image_path[] = { "./local_resources/images/heightmap.exr" };

	// Height map generated from geo-survey data
	char const* image_path[] = { "./local_resources/images/SO28sw_DTM_1m.exr" };

	// Note that when adding this image to the resource manager, we declare it to be watched (the last parameter set to true)
	// this means that any changes to the image will automatically get hot-reloaded.
	app->resource_manager.add_item( app->img_heightmap, image_info, image_path, true );

	return app;
}

// ----------------------------------------------------------------------

static bool pass_upload_mesh_data_setup( le_renderpass_o* pRp, void* user_data ) {

	auto app = static_cast<exr_decode_example_app_o*>( user_data );

	if ( app->was_mesh_uploaded ) {
		return false;
	}

	// ----------| Invariant: Mesh data was not yet uploaded

	le::RenderPass rp( pRp );

	rp
	    .useBufferResource( app->gpu_mesh->pos_handle, le::AccessFlagBits2::eTransferWrite )   //
	    .useBufferResource( app->gpu_mesh->uv_handle, le::AccessFlagBits2::eTransferWrite )    //
	    .useBufferResource( app->gpu_mesh->index_handle, le::AccessFlagBits2::eTransferWrite ) //
	    ;

	app->was_mesh_uploaded = true;

	return true;
}

// ----------------------------------------------------------------------

static void pass_upload_mesh_data_exec( le_command_buffer_encoder_o* encoder_, void* user_data ) {

	auto                app = static_cast<exr_decode_example_app_o*>( user_data );
	le::TransferEncoder encoder( encoder_ );

	// Upload mesh data
	{

		void*  gpu_buffer_data = nullptr;
		size_t num_vertices    = app->gpu_mesh->pos_count;
		size_t num_bytes       = 0;

		// Upload attribute data

		// Get a pointer to gpu mapped memory that will be uploaded to pos_handle at offset 0.
		// Then read the data from our mesh into that memory address.

		num_bytes = app->gpu_mesh->pos_num_bytes;
		if ( encoder.mapBufferMemory( app->gpu_mesh->pos_handle, 0, num_bytes, &gpu_buffer_data ) ) {
			app->mesh.readAttributeDataInto( gpu_buffer_data, num_bytes, le_mesh_api::attribute_name_t::ePosition );
		}

		num_bytes = app->gpu_mesh->uv_num_bytes;
		if ( encoder.mapBufferMemory( app->gpu_mesh->uv_handle, 0, num_bytes, &gpu_buffer_data ) ) {
			app->mesh.readAttributeDataInto( gpu_buffer_data, num_bytes, le_mesh_api::attribute_name_t::eUv );
		}

		// upload index data

		num_bytes = app->gpu_mesh->indices_num_bytes;
		if ( encoder.mapBufferMemory( app->gpu_mesh->index_handle, 0, num_bytes, &gpu_buffer_data ) ) {
			app->mesh.readIndexDataInto( gpu_buffer_data, num_bytes );
		}
	}
}

// ----------------------------------------------------------------------

static void pass_draw_exec( le_command_buffer_encoder_o* encoder_, void* user_data ) {
	auto                app = static_cast<exr_decode_example_app_o*>( user_data );
	le::GraphicsEncoder encoder{ encoder_ };

	le::Extent2D extents;
	encoder.getRenderpassExtent( &extents );

	le::Viewport viewports[ 1 ] = {
	    { 0.f, 0.f, float( extents.width ), float( extents.height ), 0.f, 1.f },
	};

	app->camera.setViewport( viewports[ 0 ] );

	// Data as it is laid out in the shader ubo
	struct MvpUbo_t {
		glm::mat4 model;
		glm::mat4 view;
		glm::mat4 projection;
	};

	// Draw main scene

	static auto psoDefaultGraphics =
	    LeGraphicsPipelineBuilder( encoder.getPipelineManager() )
	        .addShaderStage(
	            LeShaderModuleBuilder( encoder.getPipelineManager() )
	                .setShaderStage( le::ShaderStage::eVertex )
	                .setSourceFilePath( "./local_resources/shaders/isolines.vert" )
	                .build() )
	        .addShaderStage(
	            LeShaderModuleBuilder( encoder.getPipelineManager() )
	                .setShaderStage( le::ShaderStage::eFragment )
	                .setSourceFilePath( "./local_resources/shaders/isolines.frag" )
	                .build() )
	        .withInputAssemblyState()

	        .setTopology( le::PrimitiveTopology::eTriangleList )

	        .end()
	        .withRasterizationState()
	        // .setPolygonMode( le::PolygonMode::eLine )
	        .end()
	        .build();

	MvpUbo_t mvp;
	mvp.model = glm::mat4( 1.f ); // identity matrix
	mvp.model = glm::rotate( mvp.model, glm::two_pi<float>() * -0.25f, glm::vec3( 1, 0, 0 ) );
	app->camera.getViewMatrix( ( float* )( &mvp.view ) );
	app->camera.getProjectionMatrix( ( float* )( &mvp.projection ) );

	uint64_t bufferOffsets[ 2 ] = {
	    0,
	    0,
	};

	le_buffer_resource_handle buffers[] = {
	    app->gpu_mesh->pos_handle,
	    app->gpu_mesh->uv_handle,
	};

	encoder
	    .setLineWidth( 1 )
	    .bindGraphicsPipeline( psoDefaultGraphics )
	    .setArgumentData( LE_ARGUMENT_NAME( "Mvp" ), &mvp, sizeof( MvpUbo_t ) )
	    .bindVertexBuffers( 0, 2, buffers, bufferOffsets )
	    .setArgumentTexture( LE_ARGUMENT_NAME( "tex_unit_0" ), app->tex_unit_0 )
	    .bindIndexBuffer( app->gpu_mesh->index_handle, 0, app->gpu_mesh->index_type )
	    .drawIndexed( app->gpu_mesh->indices_count );
}

// ----------------------------------------------------------------------

static void exr_decode_example_app_process_ui_events( exr_decode_example_app_o* self ) {
	using namespace le_window;
	uint32_t         numEvents;
	LeUiEvent const* pEvents;

	window_i.get_ui_event_queue( self->window, &pEvents, &numEvents );
	auto const events_end = pEvents + numEvents;

	bool         wants_toggle = false;
	bool         was_resized  = false;
	le::Extent2D window_extents;

	// Key Mapping:
	// ------------
	// F11 : toggle fullscreen
	// Z   : reset camera
	// O   : toggle camera projection: orthographic/perspective
	// X   : set camera pivot distance to 0 (camera will not orbit a point, but "turn head")
	// C   : set camera pivot distance to distance to mesh centre

	for ( auto pEv = pEvents; pEv != events_end; pEv++ ) {
		auto& event = *pEv;
		switch ( event.event ) {
		case ( LeUiEvent::Type::eWindowExtent ): {
			auto& e        = event.windowExtent;
			window_extents = {
			    .width  = e.width,
			    .height = e.height,
			};
			was_resized = true;
		} break;

		case ( LeUiEvent::Type::eKey ): {
			auto& e = event.key;
			if ( e.action == LeUiEvent::ButtonAction::eRelease ) {
				if ( e.key == LeUiEvent::NamedKey::eF11 ) {
					wants_toggle ^= true;
				} else if ( e.key == LeUiEvent::NamedKey::eZ ) {
					reset_camera( self );
				} else if ( e.key == LeUiEvent::NamedKey::eO ) {
					bool is_orthographic = self->camera.getIsOrthographic();
					if ( is_orthographic == true ) {
						is_orthographic = false;
					} else {
						is_orthographic = true;
					}
					self->camera.setIsOrthographic( is_orthographic );
					std::cout << "is_orthographic: " << ( is_orthographic == true ? "true" : "false" ) << std::endl;

				} else if ( e.key == LeUiEvent::NamedKey::eX ) {
					self->cameraController.setPivotDistance( 0 );
				} else if ( e.key == LeUiEvent::NamedKey::eC ) {
					glm::mat4x4 view_matrix;
					self->camera.getViewMatrix( ( float* )( &view_matrix ) );
					glm::vec4 mesh_centre_in_world_space   = glm::vec4{ 0, 0, 0, 1 }; // given in world space
					glm::vec4 camera_origin_in_world_space = glm::inverse( view_matrix ) * glm::vec4( 0, 0, 0, 1 );
					float     distance_to_mesh_centre      = glm::distance( mesh_centre_in_world_space, camera_origin_in_world_space );
					self->cameraController.setPivotDistance( distance_to_mesh_centre );
				}
			} // if ButtonAction == eRelease

		} break;
		default:
			// do nothing
			break;
		}
	}

	if ( was_resized ) {
		self->renderer.resizeSwapchain( window_extents.width, window_extents.height );
		self->cameraController.setControlRect( 0, 0, float( window_extents.width ), float( window_extents.height ) );
	}

	self->cameraController.processEvents( self->camera, pEvents, numEvents );

	if ( wants_toggle ) {
		self->window.toggleFullscreen();
	}
}

// ----------------------------------------------------------------------

static bool exr_decode_example_app_update( exr_decode_example_app_o* self ) {

	// Polls events for all windows
	le::Window::pollEvents();

	if ( self->window.shouldClose() ) {
		return false;
	}

	exr_decode_example_app_process_ui_events( self );

	le::RenderGraph rg{};
	{
		self->resource_manager.update( rg );

		auto passUploadMeshData =
		    le::RenderPass( "upload_mesh_data", le::QueueFlagBits::eTransfer )
		        .setSetupCallback( self, pass_upload_mesh_data_setup )
		        .setExecuteCallback( self, pass_upload_mesh_data_exec );

		auto heightmap_sampler_info =
		    le::ImageSamplerInfoBuilder()
		        .withImageViewInfo()
		        .setImage( self->img_heightmap )
		        .end()
		        .withSamplerInfo()
		        .setAddressModeU( le::SamplerAddressMode::eClampToEdge )
		        .setAddressModeV( le::SamplerAddressMode::eClampToEdge )
		        .end()
		        .build();

		static le_image_resource_handle depth_buffer_image = self->renderer.createImageResourceHandle( "depth_buffer" );
		static le_image_resource_handle swapchain_image    = self->renderer.getSwapchainResource();

		auto passDraw =
		    le::RenderPass( "draw", le::QueueFlagBits::eGraphics )
		        .addColorAttachment( swapchain_image )
		        .addDepthStencilAttachment( depth_buffer_image )
		        .useBufferResource( self->gpu_mesh->uv_handle )
		        .useBufferResource( self->gpu_mesh->pos_handle )
		        .useBufferResource( self->gpu_mesh->index_handle, le::AccessFlagBits2::eIndexRead )
		        .sampleTexture( self->tex_unit_0, heightmap_sampler_info )
		        .setSampleCount( le::SampleCountFlagBits::e4 )
		        .setExecuteCallback( self, pass_draw_exec ) //
		    ;

		// Build rendergraph using the passes that we have declared above:

		rg
		    .addRenderPass( passUploadMeshData )
		    .addRenderPass( passDraw );

		// Declare buffer resources that will be used in the rendergraph
		// so that the backend knows how to allocate them once they are
		// actually used for the first time:

		rg
		    .declareResource(
		        self->gpu_mesh->pos_handle,
		        le::BufferInfoBuilder()
		            .setSize( self->gpu_mesh->pos_num_bytes )
		            .addUsageFlags( le::BufferUsageFlagBits::eVertexBuffer | le::BufferUsageFlagBits::eTransferDst )
		            .build() )
		    .declareResource(
		        self->gpu_mesh->uv_handle,
		        le::BufferInfoBuilder()
		            .setSize( self->gpu_mesh->uv_num_bytes )
		            .addUsageFlags( le::BufferUsageFlagBits::eVertexBuffer | le::BufferUsageFlagBits::eTransferDst )
		            .build() )
		    .declareResource(
		        self->gpu_mesh->index_handle,
		        le::BufferInfoBuilder()
		            .setSize( self->gpu_mesh->indices_num_bytes )
		            .addUsageFlags( le::BufferUsageFlagBits::eIndexBuffer | le::BufferUsageFlagBits::eTransferDst )
		            .build() ) //
		    ;
	}

	self->renderer.update( rg );

	self->frame_counter++;

	return true; // keep app alive
}
// ----------------------------------------------------------------------

static void reset_camera( exr_decode_example_app_o* self ) {
	le::Extent2D extents{};
	self->renderer.getSwapchainExtent( &extents.width, &extents.height );
	self->camera.setViewport( { 0, 0, float( extents.width ), float( extents.height ), 0.f, 1.f } );
	self->camera.setClipDistances( 0.0000001, 10000 );
	self->camera.setFovRadians( glm::radians( 60.f ) ); // glm::radians converts degrees to radians

	constexpr auto SHOULD_USE_VIEW_MATRIX_PRESET = false;

	if ( SHOULD_USE_VIEW_MATRIX_PRESET ) {
		glm::mat4 view_matrix =
		    glm::mat4{
		        { 0.923001, -0.198597, -0.329594, -0.000000 },
		        { 0.355046, 0.769793, 0.530437, 0.000000 },
		        { 0.148376, -0.606615, 0.781031, -0.000000 },
		        { 0.004781, -0.004931, -1158.092651, 1.000000 },
	        };
		self->camera.setViewMatrix( ( float* )&view_matrix );
	} else {
		glm::mat4 camMatrix = glm::lookAt( glm::vec3{ 0, 0, self->camera.getUnitDistance() }, glm::vec3{ 0 }, glm::vec3{ 0, 1, 0 } );
		self->camera.setViewMatrix( ( float* )( &camMatrix ) );
	}
	self->cameraController.setPivotDistance( self->camera.getUnitDistance() );
	self->camera.setIsOrthographic( true );
}

// ----------------------------------------------------------------------

static void exr_decode_example_app_destroy( exr_decode_example_app_o* self ) {
	if ( self ) {
		delete self->gpu_mesh;
	}
	delete ( self );
}

// ----------------------------------------------------------------------

LE_MODULE_REGISTER_IMPL( exr_decode_example_app, api ) {
	auto  exr_decode_example_app_api_i = static_cast<exr_decode_example_app_api*>( api );
	auto& exr_decode_example_app_i     = exr_decode_example_app_api_i->exr_decode_example_app_i;

	exr_decode_example_app_i.initialize = app_initialize;
	exr_decode_example_app_i.terminate  = app_terminate;

	exr_decode_example_app_i.create  = exr_decode_example_app_create;
	exr_decode_example_app_i.destroy = exr_decode_example_app_destroy;
	exr_decode_example_app_i.update  = exr_decode_example_app_update;
}
